// goboard.h
#ifndef Goboard_H
#define Goboard_H

#include "gotypes.h"
#include "zobrist.h"
#include "helpers.h"
#include <unordered_map>
#include <unordered_set>
#include <memory>
#include <vector>
#include <cassert>

class Move {
public:
    Point point;
    bool is_play;
    bool is_pass;
    bool is_resign;

    // Constructors are private to enforce factory methods
private:
    Move(Point p, bool pass_move, bool resign_move)
        : point(p), is_play(p.row != 0 || p.col != 0), is_pass(pass_move), is_resign(resign_move) {}

public:
    static Move Play(const Point& point) {
        return Move(point, false, false);
    }

    static Move Pass() {
        return Move(Point(), true, false);
    }

    static Move Resign() {
        return Move(Point(), false, true);
    }

    bool operator==(const Move& other) const {
        return (point == other.point) &&
               (is_pass == other.is_pass) &&
               (is_resign == other.is_resign);
    }
};

// Forward declaration
class GoString;

class GoString {
public:
    Player color;
    std::unordered_set<Point, PointHash> stones;
    std::unordered_set<Point, PointHash> liberties;

    GoString(Player c, const std::unordered_set<Point, PointHash>& s, const std::unordered_set<Point, PointHash>& l)
        : color(c), stones(s), liberties(l) {}

    GoString without_liberty(const Point& point) const {
        std::unordered_set<Point, PointHash> new_liberties = liberties;
        new_liberties.erase(point);
        return GoString(color, stones, new_liberties);
    }

    GoString with_liberty(const Point& point) const {
        std::unordered_set<Point, PointHash> new_liberties = liberties;
        new_liberties.insert(point);
        return GoString(color, stones, new_liberties);
    }

    GoString merged_with(const GoString& other) const {
        assert(other.color == color);
        std::unordered_set<Point, PointHash> combined_stones = stones;
        combined_stones.insert(other.stones.begin(), other.stones.end());

        std::unordered_set<Point, PointHash> combined_liberties = liberties;
        for (const auto& lib : other.liberties) {
            combined_liberties.insert(lib);
        }
        for (const auto& stone : stones) {
            combined_liberties.erase(stone);
        }
        for (const auto& stone : other.stones) {
            combined_liberties.erase(stone);
        }

        return GoString(color, combined_stones, combined_liberties);
    }

    size_t num_liberties() const {
        return liberties.size();
    }

    bool operator==(const GoString& other) const {
        return (color == other.color &&
                stones == other.stones &&
                liberties == other.liberties);
    }
};

class Board {
public:
    int num_rows;
    int num_cols;
    std::unordered_map<Point, std::shared_ptr<GoString>, PointHash> grid;
    uint64_t hash_code;
    Zobrist zobrist_table;

    Board(int rows, int cols)
        : num_rows(rows), num_cols(cols), hash_code(0), zobrist_table(rows, cols) {}

    bool is_on_grid(const Point& point) const {
        return (point.row >= 1 && point.row <= num_rows &&
                point.col >= 1 && point.col <= num_cols);
    }

    Player get(const Point& point) const {
        auto it = grid.find(point);
        if (it == grid.end()) return static_cast<Player>(0); // None
        return it->second->color;
    }

    std::shared_ptr<GoString> get_go_string(const Point& point) const {
        auto it = grid.find(point);
        if (it == grid.end()) return nullptr;
        return it->second;
    }

    uint64_t zobrist_hash() const {
        return hash_code;
    }

    void place_stone(Player player, const Point& point) {
        assert(is_on_grid(point));
        assert(grid.find(point) == grid.end());

        std::vector<std::shared_ptr<GoString>> adjacent_same_color;
        std::vector<std::shared_ptr<GoString>> adjacent_opposite_color;
        std::unordered_set<Point, PointHash> liberties;

        for (const auto& neighbor : point.neighbors()) {
            if (!is_on_grid(neighbor)) continue;
            auto neighbor_string = get_go_string(neighbor);
            if (neighbor_string == nullptr) {
                liberties.insert(neighbor);
            }
            else if (neighbor_string->color == player) {
                if (std::find(adjacent_same_color.begin(), adjacent_same_color.end(), neighbor_string) == adjacent_same_color.end()) {
                    adjacent_same_color.push_back(neighbor_string);
                }
            }
            else {
                if (std::find(adjacent_opposite_color.begin(), adjacent_opposite_color.end(), neighbor_string) == adjacent_opposite_color.end()) {
                    adjacent_opposite_color.push_back(neighbor_string);
                }
            }
        }

        auto new_string = std::make_shared<GoString>(player, std::unordered_set<Point, PointHash>{point}, liberties);

        for (const auto& same_color_string : adjacent_same_color) {
            new_string = std::make_shared<GoString>(new_string->merged_with(*same_color_string));
        }

        for (const auto& new_string_point : new_string->stones) {
            grid[new_string_point] = new_string;
        }

        hash_code ^= zobrist_table.get_code(point, player);

        for (const auto& other_color_string : adjacent_opposite_color) {
            auto replacement = other_color_string->without_liberty(point);
            if (replacement.num_liberties() > 0) {
                replace_string(other_color_string, replacement);
            }
            else {
                remove_string(other_color_string);
            }
        }
    }

private:
    void remove_string(const std::shared_ptr<GoString>& string) {
        for (const auto& point : string->stones) {
            for (const auto& neighbor : point.neighbors()) {
                if (!is_on_grid(neighbor)) continue;
                auto neighbor_string = get_go_string(neighbor);
                if (neighbor_string == nullptr || neighbor_string == string) continue;
                replace_string(neighbor_string, neighbor_string->with_liberty(point));
            }
            grid.erase(point);
            hash_code ^= zobrist_table.get_code(point, string->color);
        }
    }

    void replace_string(const std::shared_ptr<GoString>& old_string, const GoString& new_string) {
        auto new_shared_string = std::make_shared<GoString>(new_string);
        for (const auto& point : new_shared_string->stones) {
            grid[point] = new_shared_string;
        }
    }
};

class GameState {
public:
    Board board;
    Player next_player;
    std::shared_ptr<GameState> previous;
    Move last_move;
    std::unordered_set<std::pair<Player, uint64_t>, 
        struct pair_hash> previous_states;

    GameState(const Board& b, Player next_p, std::shared_ptr<GameState> prev, Move move)
        : board(b), next_player(next_p), previous(prev), last_move(move) {
        if (previous == nullptr) {
            // Empty set
        }
        else {
            previous_states = previous->previous_states;
            previous_states.emplace(std::make_pair(previous->next_player, previous->board.zobrist_hash()));
        }
    }

    static GameState NewGame(int board_size) {
        Board board(board_size, board_size);
        return GameState(board, Player::Black, nullptr, Move::Pass());
    }

    GameState apply_move(const Move& move) const {
        Board next_board = board;
        if (move.is_play) {
            next_board.place_stone(next_player, move.point);
        }
        return GameState(next_board, OtherPlayer(next_player), std::make_shared<GameState>(*this), move);
    }

    bool is_over() const {
        if (!last_move.is_resign && !last_move.is_pass) {
            return false;
        }
        if (last_move.is_resign) {
            return true;
        }
        if (previous == nullptr) {
            return false;
        }
        return last_move.is_pass && previous->last_move.is_pass;
    }

    bool is_black_corners(const Point& point) const {
        if (board.get(point) == Player::Black) {
            return true;
        }
        return is_point_an_eye(board, point, Player::Black);
    }

    Player get_winner() const {
        if (!last_move.is_pass || !previous || !previous->last_move.is_pass) {
            return next_player;
        }

        int black_dots = 0;
        double winner_standards = board.num_rows * board.num_cols - 3.25;

        for (int i = 1; i <= board.num_rows; ++i) {
            for (int j = 1; j <= board.num_cols; ++j) {
                Point point(i, j);
                if (is_black_corners(point)) {
                    black_dots++;
                }
            }
        }

        if (black_dots >= winner_standards) {
            return Player::Black;
        }
        else {
            return Player::White;
        }
    }

    bool is_move_self_capture(Player player, const Move& move) const {
        if (!move.is_play) return false;
        Board next_board = board;
        next_board.place_stone(player, move.point);
        auto new_string = next_board.get_go_string(move.point);
        return new_string->num_liberties() == 0;
    }

    bool does_move_ko(Player player, const Move& move) const {
        if (!move.is_play) return false;
        Board next_board = board;
        next_board.place_stone(player, move.point);
        std::pair<Player, uint64_t> next_situation = { OtherPlayer(player), next_board.zobrist_hash() };
        return previous_states.find(next_situation) != previous_states.end();
    }

    bool is_valid_move(const Move& move) const {
        if (is_over()) return false;
        if (move.is_pass || move.is_resign) return true;
        return (board.get(move.point) != Player::Black && board.get(move.point) != Player::White) &&
               !is_move_self_capture(next_player, move) &&
               !does_move_ko(next_player, move);
    }

    std::vector<Move> legal_moves() const {
        std::vector<Move> candidates;
        for (int r = 1; r <= board.num_rows; ++r) {
            for (int c = 1; c <= board.num_cols; ++c) {
                Point point(r, c);
                Move candidate = Move::Play(point);
                if (is_valid_move(candidate) && !is_point_an_eye(board, point, next_player)) {
                    candidates.push_back(candidate);
                }
            }
        }
        candidates.push_back(Move::Pass());
        candidates.push_back(Move::Resign());
        return candidates;
    }

private:
    // Custom hash for std::pair<Player, uint64_t>
    struct pair_hash {
        std::size_t operator()(const std::pair<Player, uint64_t>& p) const {
            return std::hash<int>()(static_cast<int>(p.first)) ^ (std::hash<uint64_t>()(p.second) << 1);
        }
    };
};

#endif // Goboard_H
